import pandas as pd
import ast
import os
import argparse
import json
import re
from typing import Dict, List, Tuple


class LabelProcessor:
    """标签处理器：支持动态配置、灵活选项和详细统计"""

    def __init__(self, config_path: str = None):
        """
        初始化标签处理器

        Args:
            config_path: 配置文件路径，如果为None则使用默认配置
        """
        self.config = self._load_config(config_path)
        self.option_columns = []  # 动态选项列名
        self.project_data_dir = self._get_project_data_dir()

    def _get_project_data_dir(self) -> str:
        """
        动态获取项目数据目录
        优先使用环境变量，其次使用相对路径
        """
        # 1. 检查环境变量
        env_data_dir = os.environ.get('PROJECT_DATA_DIR')
        if env_data_dir and os.path.exists(env_data_dir):
            return env_data_dir

        # 2. 使用指定的数据目录
        specified_dir = r"D:\sjysds\pythonProject1\data"
        if os.path.exists(specified_dir):
            return specified_dir

        # 3. 基于脚本位置推断数据目录
        script_dir = os.path.dirname(os.path.abspath(__file__))

        # 尝试不同可能的相对路径
        possible_paths = [
            os.path.join(script_dir, "..", "data"),  # src/../data
            os.path.join(script_dir, "..", "..", "data"),  # src/train_preprocessing/../../data
            os.path.join(script_dir, "data"),  # src/train_preprocessing/data
            os.path.join(os.getcwd(), "data")  # 当前工作目录/data
        ]

        for path in possible_paths:
            abs_path = os.path.abspath(path)
            if os.path.exists(abs_path):
                return abs_path

        # 4. 如果都找不到，创建默认目录
        default_dir = os.path.abspath(os.path.join(script_dir, "..", "data"))
        os.makedirs(default_dir, exist_ok=True)
        print(f"创建默认数据目录: {default_dir}")
        return default_dir

    def _load_config(self, config_path: str) -> Dict:
        """加载配置文件"""
        default_config = {
            "encoding": "GBK",
            "max_options": 10,
            "output_stats": True,
            "stats_output_path": "label_stats.json"  # 相对路径，会在数据目录下
        }

        # 如果提供了配置文件路径，尝试加载
        if config_path:
            abs_config_path = self.get_data_path(config_path)
            if os.path.exists(abs_config_path):
                try:
                    with open(abs_config_path, 'r', encoding='utf-8') as f:
                        user_config = json.load(f)
                    default_config.update(user_config)
                    print(f"成功加载配置文件: {abs_config_path}")
                except Exception as e:
                    print(f"加载配置文件失败，使用默认配置: {str(e)}")

        return default_config

    def get_data_path(self, relative_path: str) -> str:
        """
        获取数据目录下的文件绝对路径

        Args:
            relative_path: 相对于数据目录的路径

        Returns:
            绝对路径
        """
        # 如果已经是绝对路径，直接返回
        if os.path.isabs(relative_path):
            return relative_path

        # 构建数据目录下的绝对路径
        absolute_path = os.path.join(self.project_data_dir, relative_path)

        # 确保目录存在
        dir_path = os.path.dirname(absolute_path)
        if dir_path and not os.path.exists(dir_path):
            os.makedirs(dir_path, exist_ok=True)

        return absolute_path

    def detect_option_columns(self, df: pd.DataFrame) -> List[str]:
        """
        动态检测选项列名
        """
        option_cols = []

        # 检测opt_开头的列
        opt_pattern = re.compile(r'^opt_([A-Za-z])$')
        for col in df.columns:
            match = opt_pattern.match(col)
            if match:
                option_cols.append(match.group(1).upper())

        # 按字母顺序排序
        option_cols.sort()

        if not option_cols:
            # 如果没有检测到，使用默认的A-E
            option_cols = ['A', 'B', 'C', 'D', 'E']
            print("未检测到选项列，使用默认选项: A-E")
        else:
            print(f"检测到选项列: {option_cols}")

        return option_cols

    def get_binary_labels(self, correct_list: List[str], option_columns: List[str]) -> Dict[str, int]:
        """
        将正确选项列表转换为二进制标签字典
        """
        labels = {f'label_{opt}': 0 for opt in option_columns}

        for opt in correct_list:
            opt_upper = opt.upper().strip()
            if opt_upper in option_columns:
                labels[f'label_{opt_upper}'] = 1

        return labels

    def ensure_list_format(self, correct_val) -> List[str]:
        """
        确保correct_options为列表格式
        """
        if pd.isna(correct_val):
            return []

        if isinstance(correct_val, str):
            try:
                # 尝试解析字符串为列表
                result = ast.literal_eval(correct_val)
                if isinstance(result, list):
                    return [str(item).upper().strip() for item in result]
                else:
                    return []
            except (ValueError, SyntaxError):
                # 如果是逗号分隔的字符串，如"A,C"
                if ',' in correct_val:
                    return [opt.strip().upper() for opt in correct_val.split(',')]
                else:
                    return [correct_val.strip().upper()]
        elif isinstance(correct_val, list):
            return [str(item).upper().strip() for item in correct_val]
        else:
            return []

    def check_file_exists(self, file_path: str) -> bool:
        """
        检查文件是否存在，如果不存在则提示可能的路径
        """
        if os.path.exists(file_path):
            return True

        print(f"错误：文件未找到 - {file_path}")
        print(f"数据目录: {self.project_data_dir}")

        # 搜索数据目录中的CSV文件
        print("\n数据目录中的CSV文件:")
        try:
            for file in os.listdir(self.project_data_dir):
                if file.endswith('.csv'):
                    print(f"  - {file}")
        except Exception:
            pass

        return False

    def generate_statistics(self, df: pd.DataFrame) -> Dict:
        """
        生成详细的统计信息
        """
        total_questions = len(df)

        # 题型分布
        question_type_dist = df['correct_count'].value_counts().to_dict()

        # 单选项分布
        option_correct_dist = {}
        for opt in self.option_columns:
            col_name = f'label_{opt}'
            if col_name in df.columns:
                option_correct_dist[opt] = int(df[col_name].sum())

        # 多选项组合统计
        def get_option_combination(row):
            selected = []
            for opt in self.option_columns:
                if row[f'label_{opt}'] == 1:
                    selected.append(opt)
            return '-'.join(selected) if selected else 'None'

        combination_stats = df.apply(get_option_combination, axis=1).value_counts().to_dict()

        stats = {
            "total_questions": total_questions,
            "question_type_distribution": question_type_dist,
            "single_option_correct_distribution": option_correct_dist,
            "option_combination_statistics": combination_stats,
            "single_choice_percentage": (
                        question_type_dist.get(1, 0) / total_questions * 100) if total_questions > 0 else 0,
            "multi_choice_percentage": ((total_questions - question_type_dist.get(1,
                                                                                  0)) / total_questions * 100) if total_questions > 0 else 0,
            "data_directory": self.project_data_dir
        }

        return stats

    def save_statistics(self, stats: Dict, output_path: str = None):
        """
        保存统计信息到文件
        """
        if output_path is None:
            output_path = self.config.get('stats_output_path', 'label_stats.json')

        output_path = self.get_data_path(output_path)

        try:
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(stats, f, ensure_ascii=False, indent=2)
            print(f"统计信息已保存至: {output_path}")
        except Exception as e:
            print(f"保存统计信息失败: {str(e)}")

    def print_statistics(self, stats: Dict):
        """打印统计信息"""
        print("\n" + "=" * 60)
        print("数据统计报告")
        print("=" * 60)
        print(f"数据目录: {stats['data_directory']}")
        print(f"总题目数量: {stats['total_questions']}")
        print(f"\n题型分布:")
        for count, num in sorted(stats['question_type_distribution'].items()):
            percentage = (num / stats['total_questions'] * 100) if stats['total_questions'] > 0 else 0
            print(f"  {count}个正确选项: {num}题 ({percentage:.1f}%)")

        print(f"\n单选题比例: {stats['single_choice_percentage']:.1f}%")
        print(f"多选题比例: {stats['multi_choice_percentage']:.1f}%")

        print(f"\n各选项被选为正确答案的次数:")
        for opt, count in stats['single_option_correct_distribution'].items():
            print(f"  选项{opt}: {count}次")

    def add_correct_labels_and_count(self, processed_train_path: str, output_path: str) -> Tuple[pd.DataFrame, Dict]:
        """
        为预处理后的训练集添加二进制标签和正确选项计数
        """
        # 获取绝对路径
        absolute_input_path = self.get_data_path(processed_train_path)
        absolute_output_path = self.get_data_path(output_path)

        print(f"数据目录: {self.project_data_dir}")
        print(f"输入文件: {absolute_input_path}")
        print(f"输出文件: {absolute_output_path}")

        # 检查文件是否存在
        if not self.check_file_exists(absolute_input_path):
            raise FileNotFoundError(f"预处理训练集未找到！请确认文件是否在数据目录中")

        # 1. 加载预处理后的训练集
        try:
            train_df = pd.read_csv(absolute_input_path, encoding=self.config['encoding'])
            print(f"成功加载预处理训练集，共{len(train_df)}条记录")
        except Exception as e:
            raise Exception(f"加载训练集失败：{str(e)}")

        # 2. 动态检测选项列
        self.option_columns = self.detect_option_columns(train_df)

        # 3. 确保correct_options为列表格式
        train_df['correct_options'] = train_df['correct_options'].apply(self.ensure_list_format)

        # 4. 生成二进制标签
        label_data = []
        for correct_opts in train_df['correct_options']:
            labels = self.get_binary_labels(correct_opts, self.option_columns)
            label_data.append(labels)

        label_df = pd.DataFrame(label_data)
        train_df_with_labels = pd.concat([train_df, label_df], axis=1)

        # 5. 新增correct_count字段
        train_df_with_labels['correct_count'] = train_df_with_labels['correct_options'].apply(
            lambda x: len([opt for opt in x if opt.upper() in self.option_columns])
        )

        # 6. 数据质量校验
        valid_mask = (train_df_with_labels['correct_count'] >= 1) & (
                train_df_with_labels['correct_count'] <= len(self.option_columns))
        train_df_valid = train_df_with_labels[valid_mask].copy()
        invalid_count = len(train_df_with_labels) - len(train_df_valid)

        if invalid_count > 0:
            print(f"过滤异常数据{invalid_count}条")

        # 7. 生成统计信息
        stats = self.generate_statistics(train_df_valid)

        # 8. 保存含标签的训练集
        try:
            train_df_valid.to_csv(absolute_output_path, encoding=self.config['encoding'], index=False)
            print(f"含标签的训练集已保存至：{absolute_output_path}")
            print(f"有效训练数据量：{len(train_df_valid)}条")

        except Exception as e:
            raise Exception(f"保存含标签训练集失败：{str(e)}")

        # 9. 输出统计信息
        if self.config.get('output_stats', True):
            self.print_statistics(stats)
            self.save_statistics(stats)

        return train_df_valid, stats


def main():
    """主函数：支持命令行参数"""
    parser = argparse.ArgumentParser(description='为训练集添加二进制标签和统计信息')
    parser.add_argument('--input', '-i',
                        default='processed_train.csv',
                        help='预处理后的训练集文件名，默认: processed_train.csv')
    parser.add_argument('--output', '-o',
                        default='train_with_labels.csv',
                        help='输出文件名，默认: train_with_labels.csv')
    parser.add_argument('--config', '-c',
                        default=None,
                        help='配置文件路径，可选')
    parser.add_argument('--no-stats', action='store_true',
                        help='不输出统计信息')

    args = parser.parse_args()

    # 初始化处理器
    processor = LabelProcessor(args.config)

    if args.no_stats:
        processor.config['output_stats'] = False

    try:
        # 执行标签化处理
        train_df_final, stats = processor.add_correct_labels_and_count(args.input, args.output)

        # 打印处理结果示例
        print("\n处理结果示例（前3条）：")
        sample_cols = ['train_question_id', 'cleaned_question', 'correct_options', 'correct_count']
        sample_cols.extend([f'label_{opt}' for opt in processor.option_columns[:3]])

        print(train_df_final[sample_cols].head(3).to_string(index=False))

    except FileNotFoundError as e:
        print(f"文件未找到错误：{str(e)}")
        print("\n请确保预处理步骤已经完成，生成 processed_train.csv 文件")
        print(f"文件应该位于: {processor.project_data_dir}")
    except Exception as e:
        print(f"处理失败：{str(e)}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()